package com.shitou.springai1.wx.controller;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

public class ChatWebSocketHandler extends TextWebSocketHandler {

    private final ObjectMapper objectMapper = new ObjectMapper();
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    private final ScheduledExecutorService executorService = Executors.newScheduledThreadPool(10);

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = session.getId();
        sessions.put(sessionId, session);
        System.out.println("WebSocket 连接建立: " + sessionId);

        // 发送连接成功消息
        sendMessage(session, createResponse("连接成功", false, false));
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        System.out.println("收到消息: " + payload);

        try {
            // 解析客户端消息
            ChatRequest request = objectMapper.readValue(payload, ChatRequest.class);
            String question = request.getQuestion();
            String sessionId = request.getSessionId();

            // 处理聊天消息（模拟流式响应）
            handleChatMessage(session, question, sessionId);

        } catch (Exception e) {
            System.err.println("解析消息失败: " + e.getMessage());
            sendMessage(session, createErrorResponse("消息格式错误"));
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        System.err.println("WebSocket 传输错误: " + exception.getMessage());
        sessions.remove(session.getId());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        String sessionId = session.getId();
        sessions.remove(sessionId);
        System.out.println("WebSocket 连接关闭: " + sessionId + ", 状态: " + closeStatus);
    }

    /**
     * 处理聊天消息 - 模拟流式响应
     */
    private void handleChatMessage(WebSocketSession session, String question, String sessionId) {
        // 模拟 AI 思考的回复内容
        String fullResponse = generateAIResponse(question);

        // 模拟流式输出
        final int chunkSize = 3; // 每次发送的字符数
        final int delayMs = 100; // 每次发送的延迟

        executorService.submit(() -> {
            try {
                StringBuilder sentContent = new StringBuilder();
                String[] words = fullResponse.split("");

                for (int i = 0; i < words.length; i += chunkSize) {
                    int end = Math.min(i + chunkSize, words.length);
                    String chunk = String.join("", java.util.Arrays.copyOfRange(words, i, end));
                    sentContent.append(chunk);

                    // 发送流式数据
                    ChatResponse response = new ChatResponse();
                    response.setContent(chunk);
                    response.setStreaming(true);
                    response.setEnd(false);

                    // 在特定位置添加搜索结果的模拟数据
                    if (sentContent.length() > 20 && sentContent.length() < 30) {
                        response.setSearchResult(createMockSearchResult());
                    }

                    // 添加意图识别
                    if (sentContent.length() > 10) {
                        response.setIntent(analyzeIntent(question));
                    }

                    sendMessage(session, response);

                    // 模拟网络延迟
                    Thread.sleep(delayMs);
                }

                // 发送结束消息
                ChatResponse endResponse = new ChatResponse();
                endResponse.setContent("");
                endResponse.setStreaming(false);
                endResponse.setEnd(true);
                endResponse.setSearchResult(createMockSearchResult());
                endResponse.setIntent(analyzeIntent(question));

                sendMessage(session, endResponse);

            } catch (Exception e) {
                System.err.println("处理消息时发生错误: " + e.getMessage());
                try {
                    sendMessage(session, createErrorResponse("处理消息时发生错误"));
                } catch (IOException ex) {
                    System.err.println("发送错误消息失败: " + ex.getMessage());
                }
            }
        });
    }

    /**
     * 生成 AI 回复（模拟）
     */
    private String generateAIResponse(String question) {
        // 这里可以集成真实的 AI 服务
        return "您好！关于 \"" + question + "\" 的问题，这是我的回答：这是一个模拟的流式回复，演示了如何逐步显示AI的思考过程。在实际应用中，这里会连接真实的AI模型来生成回答。";
    }

    /**
     * 创建模拟搜索结果
     */
    private Map<String, Object> createMockSearchResult() {
        return Map.of(
            "images", java.util.List.of(
                Map.of("id", "img1", "url", "https://example.com/image1.jpg", "name", "相关图片1"),
                Map.of("id", "img2", "url", "https://example.com/image2.jpg", "name", "相关图片2")
            )
        );
    }

    /**
     * 分析意图
     */
    private String analyzeIntent(String question) {
        if (question.contains("图片") || question.contains("照片") || question.contains("图")) {
            return "image_search";
        } else if (question.contains("天气")) {
            return "weather_query";
        } else if (question.contains("时间")) {
            return "time_query";
        } else {
            return "general_chat";
        }
    }

    /**
     * 发送消息到客户端
     */
    private void sendMessage(WebSocketSession session, Object message) throws IOException {
        if (session.isOpen()) {
            String json = objectMapper.writeValueAsString(message);
            session.sendMessage(new TextMessage(json));
        }
    }

    /**
     * 创建响应对象
     */
    private ChatResponse createResponse(String content, boolean streaming, boolean isEnd) {
        ChatResponse response = new ChatResponse();
        response.setContent(content);
        response.setStreaming(streaming);
        response.setEnd(isEnd);
        return response;
    }

    /**
     * 创建错误响应
     */
    private ChatResponse createErrorResponse(String error) {
        ChatResponse response = new ChatResponse();
        response.setContent(error);
        response.setStreaming(false);
        response.setEnd(true);
        return response;
    }

    /**
     * 请求消息格式
     */
    public static class ChatRequest {
        private String question;
        private String sessionId;

        // getters and setters
        public String getQuestion() { return question; }
        public void setQuestion(String question) { this.question = question; }
        public String getSessionId() { return sessionId; }
        public void setSessionId(String sessionId) { this.sessionId = sessionId; }
    }

    /**
     * 响应消息格式
     */
    public static class ChatResponse {
        private String content;
        private boolean streaming;
        private boolean isEnd;
        private Map<String, Object> searchResult;
        private String intent;

        // getters and setters
        public String getContent() { return content; }
        public void setContent(String content) { this.content = content; }
        public boolean isStreaming() { return streaming; }
        public void setStreaming(boolean streaming) { this.streaming = streaming; }
        public boolean isEnd() { return isEnd; }
        public void setEnd(boolean end) { isEnd = end; }
        public Map<String, Object> getSearchResult() { return searchResult; }
        public void setSearchResult(Map<String, Object> searchResult) { this.searchResult = searchResult; }
        public String getIntent() { return intent; }
        public void setIntent(String intent) { this.intent = intent; }
    }
}
